
import torch
from torch import nn
from torch.autograd import Variable

import os
import numpy as np
from sklearn.model_selection import train_test_split
import cv2

# DATADIR = r'C:\DL_DATA\dl2_DataSets\zoo'
DATADIR = r'D:\_const\large_data\DL1\_many_files\zoo'

import cv2
import numpy as np

def readData(path):
    x_images = []
    y_labels = []
    for i, j in enumerate(os.listdir(path)):
        sub_path = os.path.join(path, j)
        for image_name in os.listdir(sub_path):
            image_path = os.path.join(sub_path, image_name)
            image = cv2.imread(image_path)/255
            image = cv2.resize(image, (100, 100))
            x_images.append(image)
            y_labels.append(i)
    return np.array(x_images), np.array(y_labels)

x_images, y_labels = readData(DATADIR)

x_images = Variable(torch.Tensor(x_images))
y_labels = Variable(torch.Tensor(y_labels))
x_images = torch.transpose(x_images,1,3)

x_train, x_test, y_train, y_test = train_test_split(x_images, y_labels, test_size=0.3)
print('x_train', x_train.size())
print('x_test', x_test.size())
print('y_train', y_train.size())
print('y_test', y_test.size())


class Conv_BN_LeakyReLU(nn.Module):
    def __init__(self, in_channels, out_channels, ksize, padding=0, dilation=1):
        super(Conv_BN_LeakyReLU, self).__init__()
        self.convs = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, ksize, padding=padding, dilation=dilation),
            nn.BatchNorm2d(out_channels),
            nn.LeakyReLU(0.1, inplace=True)
        )

    def forward(self, x):
        return self.convs(x)

class DarkNet_19(nn.Module):
    def __init__(self, num_classes=1000):
        print("Initializing the darknet19 network ......")

        super(DarkNet_19, self).__init__()
        # backbone network : DarkNet-19
        # output : stride = 2, c = 32
        self.conv_1 = nn.Sequential(
            Conv_BN_LeakyReLU(3, 32, 3, 1),
            nn.MaxPool2d((2 ,2), 2),
        )

        # output : stride = 4, c = 64
        self.conv_2 = nn.Sequential(
            Conv_BN_LeakyReLU(32, 64, 3, 1),
            nn.MaxPool2d((2 ,2), 2)
        )

        # output : stride = 8, c = 128
        self.conv_3 = nn.Sequential(
            Conv_BN_LeakyReLU(64, 128, 3, 1),
            Conv_BN_LeakyReLU(128, 64, 1),
            Conv_BN_LeakyReLU(64, 128, 3, 1),
            nn.MaxPool2d((2 ,2), 2)
        )

        # output : stride = 16, c = 256
        self.conv_4 = nn.Sequential(
            Conv_BN_LeakyReLU(128, 256, 3, 1),
            Conv_BN_LeakyReLU(256, 128, 1),
            Conv_BN_LeakyReLU(128, 256, 3, 1),
            nn.MaxPool2d((2 ,2), 2)
        )

        # output : stride = 32, c = 512
        self.conv_5 = nn.Sequential(
            Conv_BN_LeakyReLU(256, 512, 3, 1),
            Conv_BN_LeakyReLU(512, 256, 1),
            Conv_BN_LeakyReLU(256, 512, 3, 1),
            Conv_BN_LeakyReLU(512, 256, 1),
            Conv_BN_LeakyReLU(256, 512, 3, 1),
            nn.MaxPool2d((2 ,2), 2)
        )

        # output : stride = 32, c = 1024
        self.conv_6 = nn.Sequential(
            Conv_BN_LeakyReLU(512, 1024, 3, 1),
            Conv_BN_LeakyReLU(1024, 512, 1),
            Conv_BN_LeakyReLU(512, 1024, 3, 1),
            Conv_BN_LeakyReLU(1024, 512, 1),
            Conv_BN_LeakyReLU(512, 1024, 3, 1)
        )

        self.conv_7 = nn.Conv2d(1024, 2, 1)
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))

    def forward(self, x):
        x = self.conv_1(x)
        # print('conv_1 output: {}'.format(x.shape))
        x = self.conv_2(x)
        # print('conv_2 output: {}'.format(x.shape))
        x = self.conv_3(x)
        # print('conv_3 output: {}'.format(x.shape))
        x = self.conv_4(x)
        # print('conv_4 output: {}'.format(x.shape))
        x = self.conv_5(x)
        # print('conv_5 output: {}'.format(x.shape))
        x = self.conv_6(x)
        # print('conv_6 output: {}'.format(x.shape))

        x = self.conv_7(x)
        # print('conv_7 output: {}'.format(x.shape))
        x = self.avgpool(x)
        # print('avgpool output: {}'.format(x.shape))
        x = torch.squeeze(x)
        return x

if __name__ == '__main__':

    model = DarkNet_19()

    criterion=torch.nn.CrossEntropyLoss()

    optimizer=torch.optim.Adam(model.parameters(), lr=0.001)

    for step in range(10):
        optimizer.zero_grad()
        h = model(x_train)
        cost = criterion(h, y_train.long())
        cost.backward()
        optimizer.step()

        print(step + 1, cost.item())

    pred = torch.argmax(model(x_train), 1)
    acc = (y_train == pred).float().mean()
    print('acc', acc.item())

    pred = torch.argmax(model(x_test), 1)
    acc = (y_test == pred).float().mean()
    print('acc', acc.item())
